import torch
import torch.nn as nn


class BaseLoss(nn.Module):
    def __init__(self, config):
        super(BaseLoss, self).__init__()
        self.config = config
        self.hp = config.criterion["criterion_hp"]
        self.device = config.device

    def forward(self, y_pred, y):
        raise NotImplementedError
    
class BCE(BaseLoss):
    def __init__(self, config):
        super(BCE, self).__init__(config)
        self.criterion = nn.BCELoss(reduction="mean")

    def forward(self, y_pred, y):
        # y_pred: [bs, num_classes]
        # y: [bs, label_one_hot]

        # Calculate the loss for each element
        loss = self.criterion(y_pred.squeeze(), y.squeeze())

        return loss

class FocalBCE(BaseLoss):
    def __init__(self, config):
        super(FocalBCE, self).__init__(config)
        self.gamma = float(self.hp['gamma'])
        self.alpha = float(self.hp['alpha'])
        if self.gamma < 0 or self.alpha < 0 or self.alpha > 1:
            raise ValueError("Invalid values for gamma or alpha in Focal Loss.")
        self.criterion = nn.BCELoss(reduction="none")
    
    def forward(self, y_pred, y):
        epsilon = 1e-8
        ce_loss = self.criterion(y_pred.squeeze(), y.squeeze())

        # Compute the probability p_t
        p_t = torch.where(y == 1, y_pred, 1 - y_pred)

        # Modulating factor (1 - p_t)^gamma
        modulating_factor = torch.pow((1 - p_t + epsilon), self.gamma)

        # Apply class weighting alpha_t
        alpha_t = torch.where(y == 1, self.alpha, 1 - self.alpha)

        # Compute the focal loss
        focal_loss = -alpha_t * modulating_factor * torch.log(p_t + epsilon) * ce_loss
        focal_loss = focal_loss.squeeze()
        loss = torch.mean(focal_loss)

        return loss
    
class FocalLoss(torch.nn.Module):
    '''
        https://github.com/clcarwin/focal_loss_pytorch/blob/master/focalloss.py
    '''

    def __init__(self, gamma=0.5, alpha=None, size_average=True):
        super(FocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        if isinstance(alpha, (float, int)):
            self.alpha = torch.Tensor([alpha, 1-alpha])
        if isinstance(alpha, list):
            self.alpha = torch.Tensor(alpha)
        self.size_average = size_average

    def forward(self, input, target):
        if input.dim() > 2:
            input = input.view(input.size(0), input.size(1), -1)  # N,C,H,W => N,C,H*W
            input = input.transpose(1, 2)                         # N,C,H*W => N,H*W,C
            input = input.contiguous().view(-1, input.size(2))    # N,H*W,C => N*H*W,C
        target = target.view(-1, 1)

        logpt = F.log_softmax(input, dim=1)
        logpt = logpt.gather(1, target)
        logpt = logpt.view(-1)
        pt = torch.autograd.Variable(logpt.data.exp())

        if self.alpha is not None:
            if self.alpha.type() != input.data.type():
                self.alpha = self.alpha.type_as(input.data)
            at = self.alpha.gather(0, target.data.view(-1))
            logpt = logpt * torch.autograd.Variable(at)

        loss = -1 * (1-pt)**self.gamma * logpt
        if self.size_average:
            return loss.mean()
        else:
            return loss.sum()